package com.c2b.algorithm.leetcode.base;

import java.util.HashMap;
import java.util.Map;

/**
 * @author c2b
 * @since 2023/11/16 13:26
 */
public class LC0337 {

    static class Solution {
        public int rob(TreeNode root) {
            if (root == null) {
                return 0;
            }
            // 1.如果选择偷取当前节点的值，那么下一次不能偷取儿子节点的值，只能偷取孙子节点的值。
            int yesMoney = root.val;
            if (root.left != null) {
                yesMoney += rob(root.left.left) + rob(root.left.right);
            }
            if (root.right != null) {
                yesMoney += rob(root.right.left) + rob(root.right.right);
            }
            // 2.如果选择不偷当前节点的值，那么下一次可以选择是否偷取儿子节点的值
            int noMoney = rob(root.left) + rob(root.right);
            return Math.max(yesMoney, noMoney);
        }

        public int rob2(TreeNode root) {
            Map<TreeNode, Integer> map = new HashMap<>();
            return robInternal(root, map);
        }

        public int robInternal(TreeNode root, Map<TreeNode, Integer> map) {
            if (root == null) {
                return 0;
            }
            if (map.containsKey(root)) {
                return map.get(root);
            }
            // 1.如果选择偷取当前节点的值，那么下一次不能偷取儿子节点的值，只能偷取孙子节点的值。
            int money = root.val;
            if (root.left != null) {
                money += robInternal(root.left.left, map) + robInternal(root.left.right, map);
            }
            if (root.right != null) {
                money += robInternal(root.right.left, map) + robInternal(root.right.right, map);
            }
            // 2.如果选择不偷当前节点的值，那么下一次可以选择是否偷取儿子节点的值
            int result = Math.max(money, robInternal(root.left, map) + robInternal(root.right, map));
            // ==> 将每个节点的选择偷或者不偷的最大值，存到map中，方便使用时直接获取
            map.put(root, result);
            return result;
        }


        public int rob3(TreeNode root) {
            int[] result = robInternal(root);
            return Math.max(result[0], result[1]);
        }

        /**
         * 用于计算偷取当前节点的值，和不偷取当前节点的值，放入 int[]中。<br>
         * 其中 0 下标代表，不偷取当前节点的值；1 下标代表，偷取当前节点的值
         */
        public int[] robInternal(TreeNode root) {
            int[] result = new int[2];
            if (root == null) {
                return result;
            }
            int[] left = robInternal(root.left);
            int[] right = robInternal(root.right);
            result[0] = Math.max(left[0], left[1]) + Math.max(right[0], right[1]);
            result[1] = left[0] + right[0] + root.val;
            return result;
        }
    }

    public static void main(String[] args) {
        TreeNode root1 = new TreeNode(3);
        root1.left = new TreeNode(2);
        root1.right = new TreeNode(3);
        root1.left.right = new TreeNode(3);
        root1.right.right = new TreeNode(1);

        TreeNode root2 = new TreeNode(3);
        root2.left = new TreeNode(4);
        root2.right = new TreeNode(5);
        root2.left.left = new TreeNode(1);
        root2.left.right = new TreeNode(3);
        root2.right.right = new TreeNode(1);

        Solution solution = new Solution();
        System.out.println(solution.rob2(root1));
        System.out.println(solution.rob2(root2));
    }
}
